/*
 * Copyright (c) 2024 OpenHiTLS. All rights reserved.
 * Licensed under the OpenHiTLS license.
 */
package org.openhitls.tls;

import javax.net.ssl.SSLSession;
import javax.net.ssl.SSLSessionContext;
import java.util.Collections;
import java.util.Enumeration;
import java.util.LinkedHashMap;
import java.util.Map;

/**
 * SSL Session Context implementation backed by openHiTLS.
 *
 * <p>This class manages SSL session caching for session resumption.</p>
 */
public class HiTlsSSLSessionContext implements SSLSessionContext {

    private final boolean isServerContext;
    private int sessionCacheSize = 20480; // Default session cache size
    private int sessionTimeout = 86400;   // Default timeout: 24 hours (in seconds)

    // Session cache using LRU eviction
    private final Map<ByteArrayKey, SSLSession> sessionCache;

    HiTlsSSLSessionContext(boolean isServerContext) {
        this.isServerContext = isServerContext;
        this.sessionCache = Collections.synchronizedMap(
            new LinkedHashMap<ByteArrayKey, SSLSession>(16, 0.75f, true) {
                @Override
                protected boolean removeEldestEntry(Map.Entry<ByteArrayKey, SSLSession> eldest) {
                    if (sessionCacheSize <= 0) {
                        return false;
                    }
                    if (size() > sessionCacheSize) {
                        return true;
                    }
                    // Also check for session timeout
                    SSLSession session = eldest.getValue();
                    if (session != null && !session.isValid()) {
                        return true;
                    }
                    long currentTime = System.currentTimeMillis();
                    long sessionAge = currentTime - session.getCreationTime();
                    return sessionAge > (sessionTimeout * 1000L);
                }
            }
        );
    }

    @Override
    public SSLSession getSession(byte[] sessionId) {
        if (sessionId == null || sessionId.length == 0) {
            return null;
        }

        SSLSession session = sessionCache.get(new ByteArrayKey(sessionId));
        if (session != null) {
            // Validate session
            if (!session.isValid()) {
                sessionCache.remove(new ByteArrayKey(sessionId));
                return null;
            }

            // Check timeout
            long currentTime = System.currentTimeMillis();
            long sessionAge = currentTime - session.getCreationTime();
            if (sessionAge > (sessionTimeout * 1000L)) {
                session.invalidate();
                sessionCache.remove(new ByteArrayKey(sessionId));
                return null;
            }
        }

        return session;
    }

    @Override
    public Enumeration<byte[]> getIds() {
        // Return enumeration of valid session IDs
        return Collections.enumeration(
            sessionCache.keySet().stream()
                .filter(key -> {
                    SSLSession session = sessionCache.get(key);
                    return session != null && session.isValid();
                })
                .map(ByteArrayKey::getBytes)
                .toList()
        );
    }

    @Override
    public void setSessionTimeout(int seconds) throws IllegalArgumentException {
        if (seconds < 0) {
            throw new IllegalArgumentException("Session timeout cannot be negative");
        }
        this.sessionTimeout = seconds;
    }

    @Override
    public int getSessionTimeout() {
        return sessionTimeout;
    }

    @Override
    public void setSessionCacheSize(int size) throws IllegalArgumentException {
        if (size < 0) {
            throw new IllegalArgumentException("Cache size cannot be negative");
        }
        this.sessionCacheSize = size;

        // Trim cache if necessary
        if (size > 0) {
            synchronized (sessionCache) {
                while (sessionCache.size() > size) {
                    // Remove oldest entry
                    var iterator = sessionCache.entrySet().iterator();
                    if (iterator.hasNext()) {
                        iterator.next();
                        iterator.remove();
                    }
                }
            }
        }
    }

    @Override
    public int getSessionCacheSize() {
        return sessionCacheSize;
    }

    /**
     * Add a session to the cache.
     */
    void addSession(SSLSession session) {
        if (session == null) return;
        byte[] id = session.getId();
        if (id == null || id.length == 0) return;

        sessionCache.put(new ByteArrayKey(id), session);
    }

    /**
     * Remove a session from the cache.
     */
    void removeSession(SSLSession session) {
        if (session == null) return;
        byte[] id = session.getId();
        if (id == null || id.length == 0) return;

        sessionCache.remove(new ByteArrayKey(id));
    }

    /**
     * Check if this is a server-side context.
     */
    boolean isServerContext() {
        return isServerContext;
    }

    /**
     * Key wrapper for byte arrays to use in HashMap.
     */
    private static class ByteArrayKey {
        private final byte[] bytes;
        private final int hashCode;

        ByteArrayKey(byte[] bytes) {
            this.bytes = bytes.clone();
            this.hashCode = java.util.Arrays.hashCode(bytes);
        }

        byte[] getBytes() {
            return bytes.clone();
        }

        @Override
        public boolean equals(Object obj) {
            if (this == obj) return true;
            if (!(obj instanceof ByteArrayKey other)) return false;
            return java.util.Arrays.equals(bytes, other.bytes);
        }

        @Override
        public int hashCode() {
            return hashCode;
        }
    }
}
